from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from docx.shared import Cm
from docxtpl import DocxTemplate, InlineImage

# Исходные данные
a = 0.25
b = 0.25
a1 = 0.1
a2 = 0.2
b1 = 0.15
b2 = 0.2
h = 0.001
e = 0.7 * 10**11
Mu = 0.34
Rho = 2700
Pi = 3.14
q = (e * h**4) / (a * b) ** 2 * 0.0015
k = (e * h**3) / (a * b) ** 2 * 0.01
Dd = (e * h**3) / (12 * (1 - Mu**2))


# Функция W
def W(x: np.ndarray, y: np.ndarray) -> np.ndarray:
    """Определение перемещения

    Parameters
    ----------
    x : np.ndarray
        координата х
    y : np.ndarray
        координата у

    Returns
    -------
    np.ndarray
    Перемещения
    """
    result = 0
    for i in range(1, 4):
        for j in range(1, 4):
            qij = (
                (4 * q / (Pi**2 * i * j))
                * (np.cos((i * Pi * a2) / a) - np.cos((j * Pi * a1) / a))
                * (np.cos((j * Pi * b2) / b) - np.cos((i * Pi * b1) / b))
            )
            wij = qij / ((Dd * (((i * Pi) / a) ** 2 + ((i * Pi) / b) ** 2) ** 2) - k)
            result += wij * np.sin((i * Pi * x) / a) * np.sin((j * Pi * y) / b)
    return result


# Графики W
x_vals = np.linspace(0, a, 100)
y_vals = np.linspace(0, b, 100)
X, Y = np.meshgrid(x_vals, y_vals)
Z = W(X, Y)

fig = plt.figure()
ax = fig.add_subplot(121, projection="3d")
ax.plot_surface(X, Y, Z, cmap="viridis")
ax.set_xlabel("a")
ax.set_ylabel("b")
ax.set_zlabel("W(x,y)")
ax.set_title("W[x,y]")

ax = fig.add_subplot(122)
plt.contourf(X, Y, Z, cmap="viridis")
plt.colorbar()
plt.xlabel("a")
plt.ylabel("b")
plt.title("W[x,y]")


# Функция Mx
def Mx(x: np.ndarray, y: np.ndarray) -> np.ndarray:
    """Определяем изгибающие моменты в сечении, перпендикулярном оси х

    Parameters
    ----------
    x : np.ndarray
        координата х
    y : np.ndarray
        координата у

    Returns
    -------
    np.ndarray
    изгибающие моменты в сечении, перпендикулярном оси х
    """
    return -Dd * (
        np.gradient(np.gradient(W(x, y), axis=0), axis=0)
        + Mu * np.gradient(np.gradient(W(x, y), axis=1), axis=1)
    )


# Графики Mx
Z_Mx = Mx(X, Y)

fig_Mx = plt.figure()
ax_Mx = fig_Mx.add_subplot(121, projection="3d")
ax_Mx.plot_surface(X, Y, Z_Mx, cmap="viridis")
ax_Mx.set_xlabel("a")
ax_Mx.set_ylabel("b")
ax_Mx.set_zlabel("Mx(x,y)")
ax_Mx.set_title("Mx[x,y]")

ax_Mx = fig_Mx.add_subplot(122)
plt.contourf(X, Y, Z_Mx, cmap="viridis")
plt.colorbar()
plt.xlabel("a")
plt.ylabel("b")
plt.title("Mx[x,y]")


# Функция My
def My(x: np.ndarray, y: np.ndarray) -> np.ndarray:
    """Определяем изгибающие моменты в сечении, перпендикулярном оси у

    Parameters
    ----------
    x : np.ndarray
        координата х
    y : np.ndarray
        координата у

    Returns
    -------
    np.ndarray
    изгибающие моменты в сечении, перпендикулярном оси у
    """
    return -Dd * (
        np.gradient(np.gradient(W(x, y), axis=1), axis=1)
        + Mu * np.gradient(np.gradient(W(x, y), axis=0), axis=0)
    )


# Графики My
Z_My = My(X, Y)

fig_My = plt.figure()
ax_My = fig_My.add_subplot(121, projection="3d")
ax_My.plot_surface(X, Y, Z_My, cmap="viridis")
ax_My.set_xlabel("a")
ax_My.set_ylabel("b")
ax_My.set_zlabel("My(x,y)")
ax_My.set_title("My[x,y]")

ax_My = fig_My.add_subplot(122)
plt.contourf(X, Y, Z_My, cmap="viridis")
plt.colorbar()
plt.xlabel("a")
plt.ylabel("b")
plt.title("My[x,y]")


# Функция H
def H(x: np.ndarray, y: np.ndarray) -> np.ndarray:
    """Определяем крутящий момент

    x : np.ndarray
         координата х
     y : np.ndarray
         координата у


     Returns
     -------
     np.ndarray
     Крутящий момент

    """
    return -Dd * (1 - Mu) * np.gradient(np.gradient(W(x, y), axis=0), axis=1)


# Графики H
Z_H = H(X, Y)

fig_H = plt.figure()
ax_H = fig_H.add_subplot(121, projection="3d")
ax_H.plot_surface(X, Y, Z_H, cmap="viridis")
ax_H.set_xlabel("a")
ax_H.set_ylabel("b")
ax_H.set_zlabel("H(x,y)")
ax_H.set_title("H[x,y]")

ax_H = fig_H.add_subplot(122)
plt.contourf(X, Y, Z_H, cmap="viridis")
plt.colorbar()
plt.xlabel("a")
plt.ylabel("b")
plt.title("H[x,y]")


# Функция Qx
def Qx(x: np.ndarray, y: np.ndarray) -> np.ndarray:
    """Определяем  поперечные силы в направлении оси х
    Parameters
    ----------
    x : np.ndarray
        координата х
    y : np.ndarray
        координата у


    Returns
    -------
    np.ndarray
    Поперечные силы в направлении оси х

    """
    return -Dd * (
        np.gradient(np.gradient(np.gradient(W(x, y), axis=0), axis=0), axis=0)
        + np.gradient(np.gradient(np.gradient(W(x, y), axis=0), axis=1), axis=1)
    )


# График Qx

ZQx = Qx(X, Y)
fig_Qx = plt.figure()
ax_Qx = fig_Qx.add_subplot(121, projection="3d")
ax_Qx.plot_surface(X, Y, ZQx, cmap="viridis")
ax_Qx.set_xlabel("a")
ax_Qx.set_ylabel("b")
ax_Qx.set_zlabel("Qx(x,y)")
ax_Qx.set_title("Qx[x,y]")

ax_Qx = fig_Qx.add_subplot(122)
plt.contourf(X, Y, ZQx, cmap="viridis")
plt.colorbar()
plt.xlabel("a")
plt.ylabel("b")
plt.title("Qx[x,y]")


# Функция Qy
def Qy(x: np.ndarray, y: np.ndarray) -> np.ndarray:
    """Определяем  поперечные силы в направлении оси у
    Parameters
    ----------
    x : np.ndarray
        координата х
    y : np.ndarray
        координата у


    Returns
    -------
    np.ndarray
    Поперечные силы в направлении оси у
    """

    return -Dd * (
        np.gradient(np.gradient(np.gradient(W(x, y), axis=1), axis=1), axis=1)
        + np.gradient(np.gradient(np.gradient(W(x, y), axis=1), axis=0), axis=0)
    )


# График Qy
ZQy = Qy(X, Y)

fig_Qy = plt.figure()
ax_Qy = fig_Qy.add_subplot(121, projection="3d")
ax_Qy.plot_surface(X, Y, ZQy, cmap="viridis")
ax_Qy.set_xlabel("a")
ax_Qy.set_ylabel("b")
ax_Qy.set_zlabel("Qy(x,y)")
ax_Qy.set_title("Qy[x,y]")

ax_Qy = fig_Qy.add_subplot(122)
plt.contourf(X, Y, ZQy, cmap="viridis")
plt.colorbar()
plt.xlabel("a")
plt.ylabel("b")
plt.title("Qy[x,y]")


# Функция Tauyz
# Определяем касательные напряжения в плоскости yz

z = 0
tau_yz = 3 / 2 * Qy(X, Y) / h * (1 - (4 * z**2) / h**2)

Ztau_yz = tau_yz

fig_tau_yz = plt.figure()
ax_tau_yz = fig_tau_yz.add_subplot(121, projection="3d")
ax_tau_yz.plot_surface(X, Y, Ztau_yz, cmap="viridis")
ax_tau_yz.set_xlabel("a")
ax_tau_yz.set_ylabel("b")
ax_tau_yz.set_zlabel("tau_yz(x,y)")
ax_tau_yz.set_title("tau_yz[x,y]")

ax_tau_yz = fig_tau_yz.add_subplot(122)
plt.contourf(X, Y, Ztau_yz, cmap="viridis")
plt.colorbar()
plt.xlabel("a")
plt.ylabel("b")
plt.title("tau_yz[x,y]")

# Функция Tauxz
## Определяем касательные напряжения в плоскости xz

z = 0
tau_xz = 3 / 2 * Qx(X, Y) / h * (1 - (4 * z**2) / h**2)

Ztau_xz = tau_xz

fig_tau_xz = plt.figure()
ax_tau_xz = fig_tau_xz.add_subplot(121, projection="3d")
ax_tau_xz.plot_surface(X, Y, Ztau_xz, cmap="viridis")
ax_tau_xz.set_xlabel("a")
ax_tau_xz.set_ylabel("b")
ax_tau_xz.set_zlabel("tau_yz(x,y)")
ax_tau_xz.set_title("tau_yz[x,y]")

ax_tau_xz = fig_tau_xz.add_subplot(122)
plt.contourf(X, Y, Ztau_xz, cmap="viridis")
plt.colorbar()
plt.xlabel("a")
plt.ylabel("b")
plt.title("tau_xz[x,y]")


# Поиск максимумов tauxz

result_tauxz = np.max(tau_xz)
max_index_tauxz = np.unravel_index(np.argmax(tau_xz, axis=None), tau_xz.shape)
max_point_tauxz = (X[max_index_tauxz], Y[max_index_tauxz])


# Поиск максимумов tauyz
result_tauyz = np.max(tau_yz)
max_index_tauyz = np.unravel_index(np.argmax(tau_yz, axis=None), tau_yz.shape)
max_point_tauyz = (X[max_index_tauyz], Y[max_index_tauyz])

max_point_tauyzv = np.array(max_point_tauyz)


# Распределение максимальных касательных напряжений
Qymax = ZQy[max_index_tauyz[0], max_index_tauyz[1]]


def tau_yz1(z1: np.ndarray) -> np.ndarray:
    """Задаём фукцию максимальных касательных напряжений по толщине в плоскости yz

    Parameters
    ----------
    z1 : np.ndarray
        координата z


    Returns
    -------
    np.ndarray
    Распределение максимальных касательных напряжений по толщине в плоскости yz

    """
    return 3 / 2 * Qymax / h * (1 - (4 * z1**2) / h**2)


z1v = np.linspace(-h / 2, h / 2, 100)
tau_yz1v = tau_yz1(z1v)

fig_tauyz_max3 = plt.figure()
plt.plot(z1v, tau_yz1v)


Qxmax = ZQx[max_index_tauxz[0], max_index_tauxz[1]]


def tau_xz1(z1: np.ndarray) -> np.ndarray:
    """Задаём фукцию максимальных касательных напряжений по толщине в плоскости хz

    Parameters
    ----------
    z1 : np.ndarray
        координата z


    Returns
    -------
    np.ndarray
    Распределение максимальных касательных напряжений по толщине в плоскости хz
    """
    return 3 / 2 * Qxmax / h * (1 - (4 * z1**2) / h**2)


z1v = np.linspace(-h / 2, h / 2, 100)
tau_xz1v = tau_xz1(z1v)

fig_tauxz_max3 = plt.figure()
plt.plot(z1v, tau_xz1v)


# Заполнение шаблона

docx_template = DocxTemplate("dz/stroymech_plastina/Shablon.docx")
array = np.arange(1, 11)

# Названия рисунков
image_names = [
    "Wplot",
    "Mxplot",
    "Myplot",
    "Hplot",
    "Qxplot",
    "Qyplot",
    "Typlot",
    "Txplot",
    "Tmyplot",
    "Tmxplot",
]

for i in array:
    fig = plt.figure(i)
    fig.savefig(
        Path("dz/stroymech_plastina/" + image_names[i - 1] + ".png"), dpi=fig.dpi
    )

images = {
    image_name: InlineImage(
        docx_template, "dz/stroymech_plastina/" + image_name + ".png", Cm(12)
    )
    for image_name in image_names
}
context = {
    "YX": max_point_tauyz[0],
    "YY": max_point_tauyz[1],
    "XX": max_point_tauxz[0],
    "XY": max_point_tauxz[1],
}
context = context | images

docx_template.render(context)
docx_template.save("dz/stroymech_plastina/Otchyot.docx")
